﻿using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;

namespace EvoBoost.Core
{
	public class ExperimentRunner<T>
	{
		private SplitBase<T> split;

		public ExperimentRunner()
		{
			this.Booster = new Booster<T>();
		}

		public Booster<T> Booster { get; private set; }

		public SplitBase<T> Split
		{
			get { return this.split; }
			set
			{
				if (value == null) throw new ArgumentNullException("value");
				this.split = value;
			}
		}

		public Forest<T> RunTraining(ICollection<Sample<T>> samples)
		{
			if (samples == null) throw new ArgumentNullException("samples");

			return this.Booster.Boost(samples, this.Split);
		}

		public TrainingTestStatistics<T> RunTrainingAndTest(
			ICollection<Sample<T>> trainingSamples, ICollection<Sample<T>> testSamples)
		{
			if (trainingSamples == null) throw new ArgumentNullException("trainingSamples");
			if (testSamples == null) throw new ArgumentNullException("testSamples");

			// Train model
			Forest<T> model = this.RunTraining(trainingSamples);

			// Calculate errors after each iteration
			List<double> trainingErrors = new List<double>(model.Trees.Count);
			List<double> testErrors = new List<double>(model.Trees.Count);
			for (int i = 1; i <= model.Trees.Count; ++i)
			{
				if (i == model.Trees.Count)
				{
					Forest<T> subModel = model.GetSubForest(i);
					trainingErrors.Add(trainingSamples.Average(
						s => this.Booster.LossFunction.Calculate(s.Outcome, subModel.CalculateOutcome(s.Properties))));
					testErrors.Add(testSamples.Average(
						s => this.Booster.LossFunction.Calculate(s.Outcome, subModel.CalculateOutcome(s.Properties))));
				}
				else
				{
					trainingErrors.Add(0);
					testErrors.Add(0);
				}
			}

			// Return statistics
			return new TrainingTestStatistics<T>(model, trainingErrors, testErrors);
		}

		public CrossValidationStatistics RunKFoldCrossValidation(int foldCount, ICollection<Sample<T>> samples)
		{
			if (samples == null) throw new ArgumentNullException("samples");
			if (foldCount <= 0)
				throw new ArgumentOutOfRangeException("foldCount", "Fold count should be positive.");
			if (foldCount > samples.Count)
				throw new ArgumentOutOfRangeException("foldCount", "Fold count should not exceed sample count in training set.");

			// Check experiment runner state
			this.CheckState();

			// Create folds
			List<SampleSetHolder> folds = SplitDataToFolds(foldCount, samples);

			// Training with cross-validation
			List<double> lossDistribution = new List<double>(foldCount);
			foreach (SampleSetHolder fold in folds)
			{
				// Train model
				Forest<T> model = this.Booster.Boost(fold.TrainingSet, this.Split);

				// Test model
				double loss = fold.ValidationSet.Average(
					s => this.Booster.LossFunction.Calculate(s.Outcome, model.CalculateOutcome(s.Properties)));

				// Remember loss
				lossDistribution.Add(loss);
			}

			// Return result
			return new CrossValidationStatistics(lossDistribution);
		}

		private static List<SampleSetHolder> SplitDataToFolds(int foldCount, ICollection<Sample<T>> samples)
		{
			Debug.Assert(foldCount > 0);
			Debug.Assert(samples.Count != 0);

			List<SampleSetHolder> result = new List<SampleSetHolder>(foldCount);

			// Create fold indices
			int[] indices = new int[samples.Count];
			int index = 0;
			for (int i = 0; i < samples.Count; ++i)
			{
				indices[i] = index;
				index = (index + 1) % foldCount;
			}

			// Shuffle indices
			indices.Shuffle();

			// Training with cross-validation
			for (int i = 0; i < foldCount; ++i)
			{
				// Build training and validation sets
				List<Sample<T>> trainingSet = new List<Sample<T>>();
				List<Sample<T>> validationSet = new List<Sample<T>>();
				int sampleIndex = 0;
				foreach (Sample<T> sample in samples)
				{
					if (indices[sampleIndex++] == i)
						validationSet.Add(sample);
					else
						trainingSet.Add(sample);
				}

				result.Add(new SampleSetHolder(trainingSet, validationSet));
			}

			return result;
		}

		private void CheckState()
		{
			if (this.split == null) throw new InvalidOperationException("Valid split should be specified first.");
		}

		private class SampleSetHolder
		{
			public ICollection<Sample<T>> TrainingSet { get; private set; }

			public ICollection<Sample<T>> ValidationSet { get; private set; }

			public SampleSetHolder(ICollection<Sample<T>> trainingSet, ICollection<Sample<T>> validationSet)
			{
				this.TrainingSet = trainingSet;
				this.ValidationSet = validationSet;
			}
		}
	}
}
